//
//  NNController.cpp
//  symbrion
//
//  Created by Berend Weel on 3/22/13.
//  Copyright (c) 2013 Berend Weel. All rights reserved.
//

#include <cmath>
#include <fstream>
#include <boost/lexical_cast.hpp>

#include "NNController.h"
#include "ActivationFunctionSigmoid.h"
#include "ActivationFunctionTanh.h"
#include "ActivationFunctionLinear.h"
#include "Mathlib.h"

NNController::NNController(WorldModel *worldModel, RobotType *robot)  : EAController (worldModel, robot){
    organism = NULL;
    sequence = NULL;
    network = NULL;
    nextHingePosition = 0;
    prevHingePosition = 0;
    inputHingePositions.resize(3,0);
    outputHingePositions.resize(3,0);
    
    createNetwork();
}

NNController::~NNController(){
    delete organism;
    delete sequence;
}

void NNController::step() {
    
    // TODO: Extract the tranmission of shape genome to its own class
    // and make it an option.
    if(worldModel->isEgg){
        transmitGenome();
    }else{
        setLEDs(BLUE);
    }
    
    // Hinge angles are calculated by the egg and sent to the others in the organism
    receiveMessages();
    
    setHinge();
    
    // Current positions are sent to the egg
    sendMessages();
}

void NNController::setHinge(){
    if(worldModel->isEgg){
        std::vector<LIOReal> inputs(10,0);
        for(int i=0;i<3;i++){
            // There is a delay of 1 controller step between the egg and the other robots
            if(i==worldModel->robotNumber){
                inputs[i] = prevHingePosition;
            }else{
                inputs[i] = inputHingePositions[i];
            }
        }
        inputs[5] = std::sin((worldModel->time + (TIMESTEP/1000.0)) * 0.25);
        //inputs[6] = std::cos((worldModel->time + (TIMESTEP/1000.0)) * 0.25);
        
        network->step(inputs);
        
        std::vector<LIOReal> outputs = network->getOutputValues();
        
        outputHingePositions[0] = outputs[0] * 0.5;
        outputHingePositions[1] = outputs[1] * 0.5;
        outputHingePositions[2] = outputs[2] * 0.5;        
    }
    worldModel->desiredHingePosition = nextHingePosition;
}

void NNController::createNetwork(){
    network = new MultiLayeredPerceptron(INPUTLAYERSIZE,OUTPUTLAYERSIZE,HIDDENLAYERSIZE);
    network->setActivationFunction(&ActivationFunctionTanh::apply, &ActivationFunctionTanh::apply);
}

void NNController::transmitGenome(){
    if(sequence == NULL || organism == NULL){
        createShape();
    }
    
    if(worldModel->time > 1){
        std::stringstream msg;
        for(int i=0;i<sequence->bytes.size();i++){
            msg << sequence->bytes[i];
        }
        
        for(int i=0;i<robot->getNumIREmitters();i++){
            worldModel->irMsgToSend.push_back(new IRMessage(i,msg.str(),0));
        }
        
        setLEDs(GREEN);
    }
}

void NNController::createShape(){
    organism = new Shape::Organism();
    
    organism->AddNode(organism->GetRoot(), Shape::SIDE_FRONT, Shape::SIDE_BACK);
    organism->AddNode(organism->GetRoot(), Shape::SIDE_BACK, Shape::SIDE_FRONT);
    
    if(sequence == NULL){
        sequence = new Shape::Sequence();
        sequence->Serialize(*organism);
    }
}

void NNController::setLEDs(int colour){
    for(int i=0;i<robot->getNumLEDs();i++){
        worldModel->ledColours.insert(std::pair<int,int>(i,colour));
    }
}

void NNController::receiveMessages(){
    for(int i=0;i<worldModel->msgReceived.size();i++){
        std::string msg = worldModel->msgReceived[i].getDataString();
        if(worldModel->isEgg){
            if(msg.find("hingeupdate:") != std::string::npos){
                std::string update = msg.substr(msg.find(":")+1,std::string::npos);
                std::string numStr = update.substr(0,update.find(":"));
                std::string valStr = update.substr(update.find(":")+1,std::string::npos);
                // This assumes the robots are numbered 0, 1, 2, ... with the Egg being 0.
                int num = boost::lexical_cast<int>(numStr);
                inputHingePositions[num] = boost::lexical_cast<LIOReal>(valStr);
            }
            nextHingePosition = outputHingePositions[worldModel->robotNumber];
        }else{
            if(msg.find("hingeset:") != std::string::npos){
                std::string numbers = msg.substr(msg.find(":")+1,std::string::npos);
                std::string id = numbers.substr(0,numbers.find(":"));
                int robot = boost::lexical_cast<int>(id);
                if(robot == worldModel->robotNumber){
                    std::string value = numbers.substr(numbers.find(":")+1,std::string::npos);
                    nextHingePosition = boost::lexical_cast<double>(value);
                }
            }
        }
    }
}

void NNController::sendMessages(){
    if(worldModel->isEgg){
        for(int i=0;i<3;i++){
            if(i != worldModel->robotNumber){
                std::stringstream msg;
                msg << "hingeset:" << i << ":" << outputHingePositions[i];
                
                worldModel->msgToSend.push_back(new Message("+",msg.str()));
            }
        }
        prevHingePosition = worldModel->currentHingePosition;
    }else{
        std::stringstream msg;
        msg << "hingeupdate:" << worldModel->robotNumber << ":" << worldModel->currentHingePosition;
        
        worldModel->msgToSend.push_back(new Message("+",msg.str()));
    }
}

void NNController::setGenome(std::vector<LIOReal> _genome){
    weights = _genome;
    network->loadParameters(weights);
}

Individual* NNController::randomIndividual(){
    Individual *individual = new Individual();
    unsigned int size = (INPUTLAYERSIZE * HIDDENLAYERSIZE)+(HIDDENLAYERSIZE*OUTPUTLAYERSIZE);
    
    for(int i=0;i<size;i++){
        individual->genome.push_back( ((Mathlib::getRand()*2)-1) * 4);
    }
    
    individual->sigma = Mathlib::getGaussianRand(0, 1);
    
    return individual;
}

void NNController::loadFromFile(std::string file){
    if(weights.size() == 0){
        std::ifstream dmfile (file.c_str());
        while(dmfile.good()){
            double tmp;
            dmfile >> tmp;
            weights.push_back(tmp);
        }
    }
    network->loadParameters(weights);
}